# -*- coding: utf-8 -*-
"""
基础问答示例

这个示例演示了如何使用LLM后端系统进行基础的问答功能。
它展示了系统的核心组件如何协同工作，包括：

1. LLM客户端：与Ollama服务通信，生成回答
2. 混合检索器：结合关键词匹配和向量检索
3. 问答链：整合检索和生成的完整流程
4. 数据存储：管理问答数据的持久化
5. 日志系统：记录运行过程和调试信息

系统架构：
┌─────────────┐    ┌──────────────┐    ┌─────────────┐
│   用户问题   │───▶│  混合检索器   │───▶│  LLM生成    │
└─────────────┘    └──────────────┘    └─────────────┘
                           │                    │
                           ▼                    ▼
                   ┌──────────────┐    ┌─────────────┐
                   │   QA存储     │    │   最终答案   │
                   └──────────────┘    └─────────────┘

使用场景：
- 开发环境测试：验证系统功能是否正常
- 性能评估：测试不同类型问题的处理效果
- 集成测试：确保各组件协同工作
- 演示展示：向用户展示系统能力

Examples:
    运行示例：
    $ python examples/basic_qa.py
    
    预期输出：
    - 系统初始化日志
    - 每个测试问题的处理过程
    - 生成的答案内容
    - 错误处理信息（如果有）

Note:
    - 确保Ollama服务正在运行
    - 确保已安装所有依赖包
    - 检查配置文件设置是否正确
"""
import os
import sys
import traceback
from typing import List

# 添加项目根目录到Python路径
# 这样可以正确导入src包中的模块
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

# 在文件顶部的导入部分修改
from src.core.retrieval.hybrid_retriever import HybridRetriever
from src.chains.qa_chain import QAChain
from src.data.storage.postgres_qa_storage import PostgresQAStorage
from src.core.llm.ollama_client import OllamaClient  # 新增
from src.utils.logger import setup_logger
from loguru import logger
from langchain_community.embeddings import OllamaEmbeddings  # 新增ollama嵌入模型


def initialize_components():
    """
    初始化系统组件
    
    Returns:
        Tuple: (llm_client, qa_storage, retriever, qa_chain)
    """
    try:
        # 1. 初始化嵌入模型 - 改为使用Ollama
        logger.info("正在初始化嵌入模型...")
        embeddings = OllamaEmbeddings(
            model="deepseek-r1:32b",  # 或者使用其他支持嵌入的ollama模型
            base_url="http://10.48.0.81:11434"  # 使用你的ollama服务地址
        )
        logger.info("嵌入模型初始化成功")

        # 2. 初始化LLM客户端 - 改为使用OllamaClient
        logger.info("正在初始化LLM客户端...")
        llm_client = OllamaClient(
            model_name="deepseek-r1:32b",  # 或者你想使用的其他模型
            base_url="http://10.48.0.81:11434"
        )
        logger.info("LLM客户端初始化成功")

        # 3. 测试数据库连接
        logger.info("正在测试数据库连接...")
        from config.database_config import DatabaseConfig
        if not DatabaseConfig.validate_connection():
            raise Exception("数据库连接失败")
        logger.info("数据库连接成功")

        # 4. 初始化PostgreSQL存储（强制重新加载数据）
        logger.info("正在初始化PostgreSQL存储...")
        # 获取项目根目录的绝对路径
        import os
        current_dir = os.path.dirname(os.path.abspath(__file__))
        project_root = os.path.dirname(current_dir)
        data_dir = os.path.join(project_root, "data", "raw")
        
        qa_storage = PostgresQAStorage(
            embedding_function=embeddings.embed_query,
            knowledge_bases=["logic_qa.json", "flood_qa.json", "ai_qa.json"],
            data_dir=data_dir,  # 使用绝对路径
            force_reload=True  # 强制重新加载数据
        )
        logger.info("PostgreSQL存储初始化成功")

        # 5. 检查表统计信息
        logger.info("检查表统计信息...")
        table_stats = qa_storage.get_table_statistics()
        logger.info(f"表统计信息: {table_stats}")

        # 检查每个表是否有数据
        for table_name, stats in table_stats.items():
            if stats['total_count'] == 0:
                logger.warning(f"表 {table_name} 中没有数据")

        # 6. 获取所有QA数据用于混合检索器
        logger.info("正在获取QA数据...")
        qa_data = qa_storage.get_all_qa()
        logger.info(f"获取到 {len(qa_data)} 条QA数据")

        if len(qa_data) == 0:
            logger.error("没有获取到任何QA数据，请检查：")
            logger.error("1. 数据库连接是否正常")
            logger.error("2. JSON文件是否存在且格式正确")
            logger.error("3. 数据加载过程是否有错误")

            # 尝试手动检查数据文件
            import os
            import json
            missing_files = []
            
            # 获取项目根目录的绝对路径
            current_dir = os.path.dirname(os.path.abspath(__file__))
            project_root = os.path.dirname(current_dir)  # 从 examples 目录回到项目根目录
            
            for kb_file in ["logic_qa.json", "flood_qa.json", "ai_qa.json"]:
                file_path = os.path.join(project_root, "data", "raw", kb_file)
                logger.info(f"检查文件路径: {file_path}")
                if os.path.exists(file_path):
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                        logger.info(f"文件 {kb_file} 包含 {len(data)} 条记录")
                else:
                    logger.error(f"文件 {kb_file} 不存在于路径: {file_path}")
                    missing_files.append(kb_file)
            
            if missing_files:
                raise ValueError(f"缺少数据文件: {missing_files}")
            else:
                raise ValueError("数据文件存在但数据库中没有数据，请检查数据加载过程")

        # 7. 初始化混合检索器
        logger.info("正在初始化混合检索器...")
        retriever = HybridRetriever(qa_data)
        logger.info("混合检索器初始化成功")

        # 8. 初始化QA链
        logger.info("正在初始化QA链...")
        qa_chain = QAChain(llm_client, retriever)
        logger.info("QA链初始化成功")

        logger.info("所有组件初始化完成")
        return llm_client, qa_storage, retriever, qa_chain

    except Exception as e:
        traceback.print_exc()
        logger.error(f"组件初始化失败: {e}")
        raise


def get_test_questions() -> List[str]:
    """
    获取测试问题集合
    
    这些问题设计用于测试系统的不同能力：
    1. 逻辑推理问题
    2. 常识问题
    3. 创意思维问题
    4. 情境分析问题
    
    Returns:
        List[str]: 测试问题列表
        
    Note:
        这些问题来自经典的逻辑思维题，
        用于测试系统的推理和生成能力。
    """
    return [
        # "给大象装冰箱，总共分几步？",      # 经典逻辑题，测试步骤分解能力
        # "给长颈鹿装冰箱，总共分几步？",    # 变体问题，测试适应性
        # "动物园开运动会，谁没来？",        # 逻辑推理，测试关联思维
        # "你过一条河，鳄鱼住在河里，你没被咬，为什么？"  # 情境分析，测试推理能力
        """
        时间：1995年7月28日—30日，为期3天的特大暴雨过程；暴雨中心佟庄子站3天雨量532.6mm 
        最大入库洪峰流量：约10700m³/s 
        最大出库（泄洪）流量：5490m³/s；下游抚顺站观测到的峰值为6050m³/s 
        历史最高水位：136.46m（1995年7月31日04:00） 
        防洪限制水位：126.40m
        正常蓄水位：131.50m
        防洪库容：11.87亿m³
        流域面积：5437km²
        请问，这数据什么级别洪水
        """,
        "洪水最大洪峰流量是多少？",
    ]


def run_qa_test(qa_chain: QAChain, questions: List[str], logger) -> None:
    """
    运行问答测试
    
    对每个测试问题执行完整的问答流程，包括：
    1. 问题预处理
    2. 检索相关信息
    3. LLM生成回答
    4. 结果后处理
    5. 错误处理和日志记录
    
    Args:
        qa_chain (QAChain): 配置好的QA链实例
        questions (List[str]): 要测试的问题列表
        logger: 日志记录器
        
    Note:
        This function requires initialized QAChain, questions list, and logger.
        See main() function for complete usage example.
    """
    logger.info("=== 开始基础问答测试 ===")
    logger.info(f"共 {len(questions)} 个测试问题")

    success_count = 0
    total_questions = len(questions)

    for i, question in enumerate(questions, 1):
        logger.info(f"\n问题 {i}/{total_questions}: {question}")

        try:
            # 记录开始时间（用于性能分析）
            import time
            start_time = time.time()

            # 执行问答
            answer = qa_chain.run(question)

            # 计算处理时间
            processing_time = time.time() - start_time

            # 记录成功结果
            logger.info(f"答案: {answer}")
            logger.info(f"处理时间: {processing_time:.2f}秒")
            success_count += 1

        except Exception as e:
            # 详细的错误处理
            logger.error(f"问答失败: {type(e).__name__}: {e}")
            logger.debug(f"错误详情", exc_info=True)

        finally:
            # 分隔线，便于阅读
            logger.info("-" * 60)

    # 测试总结
    logger.info(f"\n=== 测试完成 ===")
    logger.info(f"成功: {success_count}/{total_questions}")
    logger.info(f"成功率: {success_count / total_questions * 100:.1f}%")


def main():
    """
    主函数 - 程序入口点
    
    执行完整的基础问答示例流程：
    1. 设置日志系统
    2. 初始化所有组件
    3. 准备测试数据
    4. 执行问答测试
    5. 处理异常和清理
    
    这个函数展示了如何正确使用系统的各个组件，
    是其他应用程序的参考实现。
    
    Raises:
        SystemExit: 当系统初始化失败时
        

    """
    # 设置日志系统
    logger = setup_logger("basic_qa_example")
    logger.info("=== 基础问答示例启动 ===")

    try:
        # 初始化系统组件
        logger.info("正在初始化系统组件...")
        llm_client, qa_storage, retriever, qa_chain = initialize_components()

        # 获取测试问题
        test_questions = get_test_questions()
        logger.info(f"准备测试 {len(test_questions)} 个问题")

        # 执行问答测试
        run_qa_test(qa_chain, test_questions, logger)

        logger.info("=== 基础问答示例完成 ===")

    except KeyboardInterrupt:
        logger.warning("用户中断程序执行")

    except Exception as e:
        logger.error(f"程序执行失败: {type(e).__name__}: {e}")
        logger.debug("详细错误信息", exc_info=True)
        sys.exit(1)

    finally:
        logger.info("程序清理完成")


if __name__ == "__main__":
    # 当作为脚本直接运行时执行主函数
    main()
